from typing import Dict, Any, List, Optional
import enum
import dataclasses

class MessageType(enum.Enum):
    """消息类型枚举"""
    INIT = "INIT"               # 初始化消息
    MODEL_UPDATE = "MODEL_UPDATE" # 模型更新
    MODEL_REQUEST = "MODEL_REQUEST" # 模型请求
    TRAIN_REQUEST = "TRAIN_REQUEST" # 训练请求
    TRAIN_COMPLETE = "TRAIN_COMPLETE" # 训练完成
    ERROR = "ERROR"             # 错误消息

@dataclasses.dataclass
class FederatedMessage:
    """联邦学习消息基类"""
    type: MessageType           # 消息类型
    sender: str                 # 发送方ID
    receiver: str               # 接收方ID
    timestamp: float            # 时间戳
    payload: Optional[Dict[str, Any]] = None # 消息负载
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            'type': self.type.value,
            'sender': self.sender,
            'receiver': self.receiver,
            'timestamp': self.timestamp,
            'payload': self.payload
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'FederatedMessage':
        """从字典创建"""
        return cls(
            type=MessageType(data['type']),
            sender=data['sender'],
            receiver=data['receiver'],
            timestamp=data['timestamp'],
            payload=data.get('payload')
        )

class FederatedLearningProtocol:
    """联邦学习协议"""
    
    @staticmethod
    def create_init_message(client_id: str) -> FederatedMessage:
        """创建初始化消息"""
        return FederatedMessage(
            type=MessageType.INIT,
            sender=client_id,
            receiver="server",
            timestamp=time.time(),
            payload={
                'capabilities': {
                    'max_batch_size': 32,
                    'supported_models': ['nn.Sequential'],
                    'compression': ['gzip', 'none'],
                    'encryption': False
                },
                'client_info': {
                    'version': '1.0',
                    'platform': platform.system()
                }
            }
        )
        
    @staticmethod
    def validate_message(msg: FederatedMessage) -> bool:
        """验证消息有效性"""
        required_fields = ['type', 'sender', 'receiver', 'timestamp']
        if not all(field in msg.to_dict() for field in required_fields):
            return False
            
        if not isinstance(msg.type, MessageType):
            return False
            
        if msg.timestamp > time.time() + 60:  # 允许1分钟时钟偏差
            return False
            
        return True
        
    @staticmethod
    def compress_message(msg: FederatedMessage) -> bytes:
        """压缩消息"""
        data = msg.to_dict()
        json_str = json.dumps(data)
        return gzip.compress(json_str.encode('utf-8'))
        
    @staticmethod
    def decompress_message(compressed: bytes) -> FederatedMessage:
        """解压缩消息"""
        json_str = gzip.decompress(compressed).decode('utf-8')
        data = json.loads(json_str)
        return FederatedMessage.from_dict(data)
    
    @staticmethod
    def create_model_request(client_id: str) -> FederatedMessage:
        """创建模型请求消息"""
        return FederatedMessage(
            type=MessageType.MODEL_REQUEST,
            sender=client_id,
            receiver="server",
            timestamp=time.time()
        )
    
    @staticmethod
    def create_model_update(
        client_id: str,
        model_state: Dict[str, Any],
        metrics: Dict[str, float]
    ) -> FederatedMessage:
        """创建模型更新消息"""
        return FederatedMessage(
            type=MessageType.MODEL_UPDATE,
            sender=client_id,
            receiver="server",
            timestamp=time.time(),
            payload={
                'model_state': model_state,
                'metrics': metrics
            }
        )
    
    @staticmethod
    def create_error_message(
        sender: str,
        receiver: str,
        error_msg: str,
        error_code: int = 500
    ) -> FederatedMessage:
        """创建错误消息"""
        return FederatedMessage(
            type=MessageType.ERROR,
            sender=sender,
            receiver=receiver,
            timestamp=time.time(),
            payload={
                'error': error_msg,
                'code': error_code
            }
        )